import pandas as pd
import info
import numpy as np
import datetime
from GetExternalDataSourceModule_CLIENT import GetExternalDataSource_CLIENT
from functools import wraps


class TradeRecord:
    def __init__(self,
                 trade_unit,
                 stop_win,
                 stop_loss,
                 raw_data_ZL,
                 raw_data_ZL_last,
                 x_min_price,
                 rotate_signal,
                 clearing_time,
                 rotate_trade_name,
                 max_trade_times,
                 long_freeze_period,
                 short_freeze_period,
                 freq
                 ):

        self.connector = GetExternalDataSource_CLIENT()  # 连接服务器
        self.calendar = self.connector.get_file('tradingCalendar', info.DATAPATH.server_root_path + '/' +
                                                'calendar' + '/' + 'tradingCalendar.pkl')
        self.calendar.index = self.calendar.TradingDate.tolist()
        self.position = {}  # 初始化交易记录
        self.count = 1  # 初始化交易笔数
        self.trade_unit = trade_unit  # 交易单位
        self.stop_win = stop_win  # 止盈条件
        self.stop_loss = stop_loss  # 止损条件
        self.raw_data = raw_data_ZL  # 回测数据
        self.raw_data_ZL_last = raw_data_ZL_last  # 上个主力合约, 合约更换日才有数据
        self.x_min_price = x_min_price  # 变动几个最小价格
        self.rotate_signal = rotate_signal  # 换合约信号
        self.clearing_time = clearing_time  # 强行平仓时间
        self.rotate_trade_name = rotate_trade_name  # 换仓价格基准
        self.contract_info = self.connector.get_file('ContractInfo',
                                                     info.DATAPATH.server_root_path + '/' + 'future_basic_data' + '/' + 'ContractInfo.pkl')  # 合约信息
        self.current_date = 0  # 当前交易日
        self.max_trade_times = max_trade_times  # 每日最大交易次数
        self.long_freeze_period = long_freeze_period  # 平多后多单冷冻时间
        self.short_freeze_period = short_freeze_period  # 平空后空单冷却时间
        self.next_long_datetime = datetime.datetime.strptime('2000-01-01 00:00:00', '%Y-%m-%d %H:%M:%S')  # 下一次多单可执行时间
        self.next_short_datetime = datetime.datetime.strptime('2000-01-01 00:00:00', '%Y-%m-%d %H:%M:%S')  # 下一次空单可执行时间
        # self.clearing_daily = True # boolean 每日收盘前是否平仓
        self.holding_status = 'Empty'  # long, short, Empty
        self.freq = freq

    def set_current_bar_trade_info(self,
                                   contract_name,
                                   close_position_price,
                                   open_position_price,
                                   last_bar_price,
                                   current_bar_datetime,
                                   current_bar_trade_date,
                                   daily_clearing_position_price=None
                                   ):
        self._contract_name = contract_name  # 合约名称
        self._close_position_price = close_position_price  # 平仓价格
        self._open_position_price = open_position_price  # 开仓价格
        self._last_bar_price = last_bar_price  # 上一个bar收盘价
        self._current_bar_datetime = current_bar_datetime  # 当前时间
        self._daily_clearing_position_price = daily_clearing_position_price  # 清仓价格
        self._current_bar_trade_date = current_bar_trade_date
        current_date_tmp = self._current_bar_trade_date  # 当前日期

        if self.current_date != current_date_tmp:
            self.new_day = True
            self.current_date = current_date_tmp
            self.trade_count_daily = 0
            self.trade_status = True
        else:
            self.new_day = False

        self.rotate_flag = self.rotate_signal[current_date_tmp]
        if len(self.clearing_time)>0:
            self._clearing_time = [(datetime.datetime.strptime(self.current_date + ' ' + x.split('_')[0], '%Y-%m-%d %H%M%S'),
                                    datetime.datetime.strptime(self.current_date + ' ' + x.split('_')[1], '%Y-%m-%d %H%M%S')
                                    ) for x in self.clearing_time]
        else:
            self._clearing_time = []
        # before generating signal check trade status

    def check_contract_rotate(self, ):
        """
        调仓日第一笔交易, 强制执行旧合约平仓
        :return:
        """
        if self.new_day and self.rotate_signal[self.current_date]:
            # 平旧仓

            self.exit_long_position_rotate()
            self.exit_short_position_rotate()

    def check_clearing_time(self):
        """
        在某个时间段强行平仓
        :return:
        """
        if len(self._clearing_time)>0:
            signal = np.sum([x[0].time()<=self._current_bar_datetime.time()<=x[1].time() for x in self._clearing_time])
            if signal:
                print('{}: meet clearing period.'.format(self._current_bar_datetime))
                if self.holding_status=='Empty':
                    pass
                elif self.holding_status=='long':
                    self.exit_long_position()
                elif self.holding_status=='short':
                    self.exit_short_position()
                self.trade_status=False
            else:
                self.trade_status=True

    def check_stop_win(self):
        """
        达到止盈条件, 强行平仓
        :return:
        """
        if self.holding_status!='Empty':
            latest_trade_num = list(self.position.keys())[-1]
            init_price = self.position[latest_trade_num]['init_price']
            latest_price = self._last_bar_price

            if self.position[latest_trade_num]['position']=='long':
                ret = latest_price/init_price - 1
            elif self.position[latest_trade_num]['position']=='short':
                ret = 1 - latest_price / init_price

            if ret>self.stop_win:
                if self.position[latest_trade_num]['position'] == 'long':
                    self.exit_long_position()
                elif self.position[latest_trade_num]['position'] == 'short':
                    self.exit_short_position()
                print('{}: stop win'.format(self._current_bar_datetime))
            else:
                pass

    def check_stop_loss(self):
        """
        达到止损条件, 强行平仓
        :return:
        """
        if self.holding_status != 'Empty':
            latest_trade_num = list(self.position.keys())[-1]
            init_price = self.position[latest_trade_num]['init_price']
            latest_price = self._last_bar_price
            if self.position[latest_trade_num]['position'] == 'long':
                ret = latest_price / init_price - 1
            elif self.position[latest_trade_num]['position'] == 'short':
                ret = 1 - latest_price / init_price

            if ret < self.stop_loss:
                print('{}: stop loss.'.format(self._current_bar_datetime))
                if self.position[latest_trade_num]['position'] == 'long':
                    self.exit_long_position()
                elif self.position[latest_trade_num]['position'] == 'short':
                    self.exit_short_position()
                print('{}: stop loss'.format(self._current_bar_datetime))
            else:
                pass

    def check_trade_num(self):
        if self.trade_count_daily >= self.max_trade_times:
            print('{}: meet max trade num limit.'.format(self._current_bar_datetime))
            self.trade_status=False
        else:
            pass

    def update_current_position(self, ):
        # check_contract_rotate_check

        self.check_contract_rotate()
        # check clearing time
        self.check_clearing_time()
        # check stop win
        self.check_stop_win()
        # check stop loss
        self.check_stop_loss()
        if self.trade_status:
            # check trade num
            self.check_trade_num()



    def enter_long_position(self):
        if (self.trade_status) and (self._current_bar_datetime>=self.next_long_datetime):
            if self.holding_status == 'long':
                pass
            elif self.holding_status=='Empty':
                trade_info_dict = {
                    'contract_name': self._contract_name,
                    'trade_unit': self.trade_unit,
                    'init_price': self._open_position_price,
                    'end_price': None,
                    'datetime_start': self._current_bar_datetime,
                    'datetime_end': None,
                    'position': 'long',
                    'ret': 0,
                    'trade_date': self.current_date
                }
                self.position[self.count] = trade_info_dict
                self.holding_status = 'long'
                self.count += 1
                self.trade_count_daily += 1
                print('{}: {} enter long position @ {}'.format(self._current_bar_datetime, self._contract_name,
                                                               self._open_position_price))
            elif self.holding_status=='short':
                trade_info_dict = {
                    'contract_name': self._contract_name,
                    'trade_unit': self.trade_unit,
                    'init_price': self._open_position_price,
                    'end_price': None,
                    'datetime_start': self._current_bar_datetime,
                    'datetime_end': None,
                    'position': 'long',
                    'ret': 0,
                    'trade_date': self.current_date
                }
                self.exit_short_position()
                self.position[self.count] = trade_info_dict
                self.holding_status = 'long'
                self.count += 1
                self.trade_count_daily += 1
                print('{}: {} enter long position @ {}'.format(self._current_bar_datetime, self._contract_name,
                                                                self._open_position_price))

    def enter_short_position(self):
        if (self.trade_status) and (self._current_bar_datetime>=self.next_short_datetime):
            if self.holding_status == 'short':
                pass
            elif self.holding_status=='Empty':
                trade_info_dict = {
                    'contract_name': self._contract_name,
                    'trade_unit': self.trade_unit,
                    'init_price': self._open_position_price,
                    'end_price': None,
                    'datetime_start': self._current_bar_datetime,
                    'datetime_end': None,
                    'position': 'short',
                    'ret': 0,
                    'trade_date': self.current_date
                }
                self.position[self.count] = trade_info_dict
                self.holding_status = 'short'
                self.count += 1
                self.trade_count_daily += 1
                print('{}: {} enter short position @ {}'.format(self._current_bar_datetime, self._contract_name,
                                                                self._open_position_price))
            elif self.holding_status=='long':
                trade_info_dict = {
                    'contract_name': self._contract_name,
                    'trade_unit': self.trade_unit,
                    'init_price': self._open_position_price,
                    'end_price': None,
                    'datetime_start': self._current_bar_datetime,
                    'datetime_end': None,
                    'position': 'short',
                    'ret': 0,
                    'trade_date': self.current_date
                }
                self.exit_long_position()
                self.position[self.count] = trade_info_dict
                self.holding_status = 'short'
                self.count += 1
                self.trade_count_daily += 1
                print('{}: {} enter short position @ {}'.format(self._current_bar_datetime, self._contract_name,
                                                                self._open_position_price))

    def exit_long_position(self,):
        if self.holding_status == 'long':
            latest_trade_num = list(self.position.keys())[-1]
            if self.position[latest_trade_num]['datetime_start'] != self._current_bar_datetime:
                self.position[latest_trade_num]['end_price'] = self._close_position_price
                self.position[latest_trade_num]['datetime_end'] = self._current_bar_datetime
                self.holding_status = 'Empty'
                self.next_long_datetime = self._current_bar_datetime + datetime.timedelta(
                    seconds=self.long_freeze_period)
                print('{}: {} exit long position @ {}'.format(self._current_bar_datetime, self._contract_name,
                                                               self._close_position_price))

        else:
            pass


    def exit_short_position(self,):
        if self.holding_status == 'short':
            latest_trade_num = list(self.position.keys())[-1]
            if self.position[latest_trade_num]['datetime_start'] != self._current_bar_datetime:
                self.position[latest_trade_num]['end_price'] = self._close_position_price
                self.position[latest_trade_num]['datetime_end'] = self._current_bar_datetime
                self.holding_status = 'Empty'
                self.next_short_datetime = self._current_bar_datetime + datetime.timedelta(
                    seconds=self.long_freeze_period)
                print('{}: {} exit short position @ {}'.format(self._current_bar_datetime, self._contract_name,
                                                              self._close_position_price))
        else:
            pass

    def exit_long_position_rotate(self,):
        if self.holding_status == 'long':
            close_position_price = self.raw_data_ZL_last.loc[self._current_bar_datetime, self.rotate_trade_name]
            latest_trade_num = list(self.position.keys())[-1]
            if self.position[latest_trade_num]['datetime_start'] != self._current_bar_datetime:
                self.position[latest_trade_num]['end_price'] = close_position_price
                self.position[latest_trade_num]['datetime_end'] = self._current_bar_datetime
                self.holding_status = 'Empty'
                print('{}: {} exit old position @ {}'.format(self._current_bar_datetime, self._contract_name,
                                                               close_position_price))
        else:
            pass


    def exit_short_position_rotate(self,):
        if self.holding_status == 'short':
            close_position_price = self.raw_data_ZL_last.loc[self._current_bar_datetime, self.rotate_trade_name]
            latest_trade_num = list(self.position.keys())[-1]
            if self.position[latest_trade_num]['datetime_start'] != self._current_bar_datetime:
                self.position[latest_trade_num]['end_price'] = close_position_price
                self.position[latest_trade_num]['datetime_end'] = self._current_bar_datetime
                self.holding_status = 'Empty'
                print('{}: {} exit old position @ {}'.format(self._current_bar_datetime, self._contract_name,
                                                              close_position_price))
        else:
            pass

    def exit_long_position_clearing_daily(self, price):
        if self.holding_status == 'long':
            latest_trade_num = list(self.position.keys())[-1]
            self.position[latest_trade_num]['end_price'] = price
            self.position[latest_trade_num]['datetime_end'] = self._current_bar_datetime
            self.holding_status = 'Empty'
            print('{}: {} exit long position @ {}'.format(self._current_bar_datetime, self._contract_name,
                                                           price))
        else:
            pass


    def exit_short_position_clearing_daily(self, price):
        if self.holding_status == 'short':
            latest_trade_num = list(self.position.keys())[-1]
            self.position[latest_trade_num]['end_price'] = price
            self.position[latest_trade_num]['datetime_end'] = self._current_bar_datetime
            self.holding_status = 'Empty'
            print('{}: {} exit short position @ {}'.format(self._current_bar_datetime, self._contract_name,
                                                          price))
        else:
            pass


    def get_final_position(self, end_date):
        final_position = pd.DataFrame(self.position).T
        final_position = final_position.merge(self.contract_info, left_on='contract_name', right_on='代码')

        # final_position['x_min_price'] = self.x_min_price
        final_position['real_ret'] = np.NaN
        final_position['real_init_price'] = np.NaN
        final_position['real_end_price'] = np.NaN
        final_position.loc[final_position.end_price.isnull(), 'end_price'] = self.raw_data.loc[self._current_bar_datetime, 'syl1']
        for index, row in final_position.iterrows():
            if row['position'] == 'short':
                real_init_price = row['init_price'] - self.x_min_price * row['最小变动价位']
                real_end_price = row['end_price'] + self.x_min_price * row['最小变动价位']
                real_init_amount = real_init_price * row['trade_unit'] * row['合约乘数']
                real_end_amount = real_end_price * row['trade_unit'] * row['合约乘数']
                real_init_amount_fee = real_init_amount * row['开仓手续费(交易额万分之)']/10000
                real_end_amount_fee = real_end_amount * row['平仓手续费(交易额万分之)']/10000
                real_init_num_fee = row['trade_unit'] * row['开仓手续费(元/手)']
                real_end_num_fee = row['trade_unit'] * row['平仓手续费(元/手)']
                final_position.loc[index, 'real_init_price'] = (real_init_amount - real_init_amount_fee - real_init_num_fee)/(row['trade_unit'] * row['合约乘数'])
                final_position.loc[index, 'real_end_price'] = (
                                                                           real_end_amount + real_end_amount_fee + real_end_num_fee) / \
                                                              (row['trade_unit'] * row['合约乘数'])
                final_position.loc[index, 'real_ret'] = (real_init_amount - real_end_amount - real_init_amount_fee - real_end_amount_fee - real_init_num_fee - real_end_num_fee)/real_init_amount
                final_position.loc[index, 'ret'] = 1 - final_position.loc[index, 'end_price']/final_position.loc[index, 'init_price']
            elif row['position'] == 'long':
                real_init_price = row['init_price'] + self.x_min_price * row['最小变动价位']
                real_end_price = row['end_price'] - self.x_min_price * row['最小变动价位']
                real_init_amount = real_init_price * row['trade_unit'] * row['合约乘数']
                real_end_amount = real_end_price * row['trade_unit'] * row['合约乘数']
                real_init_amount_fee = real_init_amount * row['开仓手续费(交易额万分之)']/10000
                real_end_amount_fee = real_end_amount * row['平仓手续费(交易额万分之)']/10000
                real_init_num_fee = row['trade_unit'] * row['开仓手续费(元/手)']
                real_end_num_fee = row['trade_unit'] * row['平仓手续费(元/手)']
                final_position.loc[index, 'real_init_price'] = (
                                                                    real_init_amount + real_init_amount_fee + real_init_num_fee) / \
                                                               (row['trade_unit'] * row['合约乘数'])
                final_position.loc[index, 'real_end_price'] = (
                                                                    real_end_amount - real_end_amount_fee - real_end_num_fee) / \
                                                              (row['trade_unit'] * row['合约乘数'])
                final_position.loc[index, 'real_ret'] = (real_end_amount - real_init_amount - real_init_amount_fee - real_end_amount_fee - real_init_num_fee - real_end_num_fee) / real_init_amount
                final_position.loc[index, 'ret'] = final_position.loc[index, 'end_price'] / final_position.loc[
                    index, 'init_price'] - 1
        final_position.index = final_position.datetime_start.tolist()
        end_data_datetime = datetime.datetime.strptime(end_date + ' 15:00:00', '%Y-%m-%d %H:%M:%S')
        final_position['status'] = 'Off'
        final_position.loc[final_position.datetime_end.isnull(), 'status'] = 'On'
        final_position.loc[final_position.datetime_end.isnull(), 'datetime_end'] = end_data_datetime

        return final_position




